#!/bin/bash

# MobileViT 树莓派部署自动安装脚本
# 使用方法: chmod +x install.sh && ./install.sh

set -e  # 遇到错误立即退出

echo "=========================================="
echo "MobileViT 树莓派部署自动安装脚本"
echo "=========================================="

# 检查是否为树莓派系统
if ! grep -q "Raspberry Pi" /proc/cpuinfo; then
    echo "警告: 此脚本专为树莓派设计，当前系统可能不兼容"
    read -p "是否继续安装? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

# 更新系统
echo "正在更新系统..."
sudo apt update && sudo apt upgrade -y

# 安装系统依赖
echo "正在安装系统依赖..."
sudo apt install -y \
    python3-pip \
    python3-venv \
    python3-dev \
    git \
    libopencv-dev \
    python3-opencv \
    espeak \
    espeak-data \
    alsa-utils \
    libatlas-base-dev \
    libopenblas-dev \
    libjpeg-dev \
    libpng-dev \
    libtiff-dev \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    libv4l-dev \
    libxvidcore-dev \
    libx264-dev \
    libgtk-3-dev \
    libcanberra-gtk-module \
    libcanberra-gtk3-module

# 创建虚拟环境
echo "正在创建Python虚拟环境..."
python3 -m venv mobilevit_env

# 激活虚拟环境
source mobilevit_env/bin/activate

# 升级pip
echo "正在升级pip..."
pip install --upgrade pip

# 安装PyTorch (CPU版本)
echo "正在安装PyTorch..."
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

# 安装其他Python依赖
echo "正在安装Python依赖..."
pip install -r requirements.txt

# 询问是否安装树莓派摄像头支持
read -p "是否安装树莓派摄像头支持? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    echo "正在安装picamera..."
    pip install picamera
    
    # 启用摄像头接口
    echo "请手动启用摄像头接口:"
    echo "sudo raspi-config -> Interface Options -> Camera -> Enable"
fi

# 配置音频
echo "正在配置音频..."
sudo usermod -a -G audio $USER

# 创建必要的目录
echo "正在创建目录结构..."
mkdir -p audio
mkdir -p logs
mkdir -p models

# 设置权限
echo "正在设置权限..."
sudo usermod -a -G video $USER
sudo usermod -a -G audio $USER

# 测试安装
echo "正在测试安装..."
python3 -c "
import torch
import cv2
import numpy as np
print('✓ PyTorch版本:', torch.__version__)
print('✓ OpenCV版本:', cv2.__version__)
print('✓ NumPy版本:', np.__version__)
print('✓ 所有依赖安装成功!')
"

# 创建启动脚本
echo "正在创建启动脚本..."
cat > start_mobilevit.sh << 'EOF'
#!/bin/bash
cd "$(dirname "$0")"
source mobilevit_env/bin/activate
python3 mobilevit_inference.py --model best_mobilevit_model.pth "$@"
EOF

chmod +x start_mobilevit.sh

# 创建系统服务文件模板
echo "正在创建系统服务模板..."
cat > mobilevit.service.template << EOF
[Unit]
Description=MobileViT Obstacle Detection
After=network.target

[Service]
Type=simple
User=$USER
WorkingDirectory=$(pwd)
Environment=PATH=$(pwd)/mobilevit_env/bin
ExecStart=$(pwd)/mobilevit_env/bin/python3 mobilevit_inference.py --model best_mobilevit_model.pth
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF

echo "=========================================="
echo "安装完成!"
echo "=========================================="
echo
echo "下一步操作:"
echo "1. 将训练好的模型文件复制到当前目录:"
echo "   cp ../best_mobilevit_model.pth ./"
echo
echo "2. 测试摄像头:"
echo "   source mobilevit_env/bin/activate"
echo "   python3 camera_capture.py"
echo
echo "3. 测试语音功能:"
echo "   python3 voice_alert.py"
echo
echo "4. 运行完整系统:"
echo "   ./start_mobilevit.sh"
echo
echo "5. 设置开机自启动 (可选):"
echo "   sudo cp mobilevit.service.template /etc/systemd/system/mobilevit.service"
echo "   sudo systemctl enable mobilevit.service"
echo "   sudo systemctl start mobilevit.service"
echo
echo "如需重启系统以应用所有配置:"
echo "   sudo reboot"
echo
echo "安装日志已保存到: install.log"

# 保存安装日志
echo "安装完成时间: $(date)" > install.log
echo "Python版本: $(python3 --version)" >> install.log
echo "系统信息: $(uname -a)" >> install.log
echo "已安装的Python包:" >> install.log
pip list >> install.log

echo "安装脚本执行完毕!"